# src/data_generation.py
import pandas as pd
import numpy as np
from faker import Faker
from datetime import datetime, timedelta
import sqlite3
import os

class DataGenerator:
    def __init__(self, seed=42):
        self.fake = Faker('zh_CN')
        np.random.seed(seed)
        Faker.seed(seed)
    
    def generate_sales_data(self, n_records=1500):
        """生成销售数据"""
        print("正在生成销售数据...")
        
        regions = ['华北', '华南', '东北', '西北', '西南', '华东']
        region_weights = [0.18, 0.12, 0.22, 0.15, 0.16, 0.17]
        
        data = []
        for i in range(n_records):
            region = np.random.choice(regions, p=region_weights)
            
            if region == '东北':
                amount_range = (2000, 4000)
            elif region == '华南':
                amount_range = (800, 2000)
            else:
                amount_range = (1000, 3000)
                
            record = {
                'order_id': f'SO{20230000 + i}',
                'customer_name': self.fake.company(),
                'product_category': np.random.choice(['电子产品', '办公用品', '服装', '食品']),
                'product_name': self.fake.word(),
                'quantity': np.random.randint(1, 20),
                'unit_price': round(np.random.uniform(50, 500), 2),
                'total_amount': round(np.random.uniform(*amount_range), 2),
                'order_date': self.fake.date_between(start_date='-1y', end_date='today'),
                'region': region,
                'sales_person': self.fake.name()
            }
            data.append(record)
        
        df = pd.DataFrame(data)
        print(f"销售数据生成完成：{len(df)} 条记录")
        return df
    
    def generate_employee_data(self, n_records=1200):
        """生成员工数据"""
        print("正在生成员工数据...")
        
        departments = ['财务部', '人力资源部', '行政部', '信息技术部', '研发部', '销售部']
        dept_sizes = [46, 9, 162, 437, 799, 300]
        
        dept_employees = []
        total_size = sum(dept_sizes)
        for i, dept in enumerate(departments):
            dept_count = int(n_records * dept_sizes[i] / total_size)
            dept_employees.extend([dept] * dept_count)
        
        while len(dept_employees) < n_records:
            dept_employees.append(np.random.choice(departments))
        
        data = []
        age_groups = {
            '20-30': (20, 30, 0.375),
            '30-40': (30, 40, 0.292),
            '40-50': (40, 50, 0.208),
            '>=50': (50, 65, 0.125)
        }
        
        education_levels = ['专科以下', '本科', '硕士研究生', '博士研究生']
        edu_weights = [0.2, 0.5, 0.25, 0.05]
        
        marital_status = ['单身', '已婚', '离异']
        marital_weights = [0.32, 0.46, 0.22]
        
        for i in range(n_records):
            age_group = np.random.choice(list(age_groups.keys()), 
                                       p=[v[2] for v in age_groups.values()])
            min_age, max_age, _ = age_groups[age_group]
            age = np.random.randint(min_age, max_age + 1)
            
            gender = np.random.choice(['男', '女'], p=[0.6, 0.4])
            
            record = {
                'employee_id': f'EMP{10000 + i}',
                'name': self.fake.name(),
                'age': age,
                'gender': gender,
                'department': dept_employees[i],
                'education': np.random.choice(education_levels, p=edu_weights),
                'marital_status': np.random.choice(marital_status, p=marital_weights),
                'hire_date': self.fake.date_between(start_date='-10y', end_date='today'),
                'salary': round(np.random.uniform(8000, 35000), 2),
                'position': self.fake.job()
            }
            data.append(record)
        
        df = pd.DataFrame(data)
        print(f"员工数据生成完成：{len(df)} 条记录")
        return df
    
    def save_data(self, df_sales, df_users):
        """保存数据"""
        print("正在保存数据文件...")
        
        df_sales.to_csv('data/raw/sales_data.csv', index=False, encoding='utf-8-sig')
        df_users.to_csv('data/raw/user_behavior.csv', index=False, encoding='utf-8-sig')
        
        conn = sqlite3.connect('data/processed/business_data.db')
        df_sales.to_sql('sales_data', conn, if_exists='replace', index=False)
        df_users.to_sql('user_behavior', conn, if_exists='replace', index=False)
        conn.close()
        
        print("数据保存完成！")

def main():
    """主函数"""
    print("开始生成模拟数据...")
    
    generator = DataGenerator()
    sales_data = generator.generate_sales_data(1500)
    user_behavior = generator.generate_employee_data(1200)
    generator.save_data(sales_data, user_behavior)
    
    print("\n数据生成流程完成！")
    return sales_data, user_behavior

if __name__ == "__main__":
    sales_df, employee_df = main()
    print("数据已生成！")